import torch
import cv2 as cv
from model import ShuffleNetV2
from utils import letter_box

model = ShuffleNetV2()
model.load_state_dict(torch.load('shufflenet.pth'))
with open('classes.txt', 'r') as r:
    classes = r.read()
image = cv.imread('rose.jpeg')
assert image is not None, 'Image does not exit.'
image = letter_box(image, 224)
image = image[:, :, ::-1].transpose(2, 1, 0)/256
print('Network loading complete.')
model.eval()
with torch.no_grad():
    image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
    predict = model(image)
print(predict)
